import fitz  # PyMuPDF
import requests
import json

# --- 1. 定义大模型服务配置 ---
# 这里使用你提供的服务信息
MODEL_SERVICES = [
    {
        "api_url": "http://18.0.31.1:9981/v1/chat/completions",
        "model": "ds671"
    }
    # 如果有备用模型，可以取消注释并添加
    # {
    #     "api_url": "http://18.0.31.2:9981/v1/chat/completions",
    #     "model": "/home/llm/qwq32"
    # }
]


def get_sensitive_info_from_llm(text):
    """
    调用大模型服务，识别并提取敏感信息。

    Args:
        text (str): 待分析的文本内容。

    Returns:
        dict: 包含敏感信息的字典，例如 {'names': ['张三'], 'id_numbers': ['...'], 'phone_numbers': ['...']}
              如果失败则返回空字典。
    """
    if not MODEL_SERVICES:
        print("未配置任何大模型服务，无法进行敏感信息识别。")
        return {}

    # 默认使用列表中的第一个模型服务
    model_config = MODEL_SERVICES[0]
    api_url = model_config["api_url"]
    model_name = model_config["model"]

    # 针对中文环境的提示词，要求返回 JSON 格式
    prompt = f"""
    请从以下文本中提取所有姓名、身份证号和手机号。
    如果找到，请以 JSON 格式返回，例如：
    {{
      "names": ["姓名1", "姓名2"],
      "id_numbers": ["身份证号1", "身份证号2"],
      "phone_numbers": ["手机号1", "手机号2"]
    }}
    如果未找到任何信息，则返回空 JSON 对象 {{}}。

    文本内容：
    {text}
    """

    headers = {
        "Content-Type": "application/json"
    }

    payload = {
        "model": model_name,
        "messages": [{"role": "user", "content": prompt}],
        "response_format": {"type": "json_object"}
    }

    try:
        response = requests.post(api_url, headers=headers, data=json.dumps(payload), timeout=60000)  # 设置超时时间
        response.raise_for_status()  # 检查 HTTP 错误

        result_json_str = response.json()["choices"][0]["message"]["content"]
        sensitive_info = json.loads(result_json_str)

        return sensitive_info

    except requests.exceptions.RequestException as e:
        print(f"调用大模型服务时发生网络错误或超时: {e}")
        return {}
    except (json.JSONDecodeError, KeyError) as e:
        print(f"解析大模型返回内容时发生错误: {e}")
        try:
            print(f"原始返回内容: {response.json()}")
        except:
            print("无法获取原始返回内容。")
        return {}


# --- 2. 整合脱敏逻辑 ---
def redact_pdf_with_llm(input_pdf_path, output_pdf_path):
    """
    结合大模型，对PDF中的姓名、身份证号和手机号进行脱敏处理。

    Args:
        input_pdf_path (str): 原始PDF文件路径。
        output_pdf_path (str): 脱敏后的PDF文件保存路径。
    """
    try:
        doc = fitz.open(input_pdf_path)
        all_sensitive_info = {}

        # 遍历PDF的每一页，提取文本并交给大模型处理
        full_text_content = ""
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            full_text_content += page.get_text() + "\n"

        print("正在调用大模型服务识别敏感信息...")
        all_sensitive_info = get_sensitive_info_from_llm(full_text_content)

        sensitive_keywords = []
        if 'names' in all_sensitive_info:
            sensitive_keywords.extend(all_sensitive_info['names'])
        if 'id_numbers' in all_sensitive_info:
            sensitive_keywords.extend(all_sensitive_info['id_numbers'])
        if 'phone_numbers' in all_sensitive_info:
            sensitive_keywords.extend(all_sensitive_info['phone_numbers'])

        if not sensitive_keywords:
            print("未识别到任何敏感信息，无需脱敏。")
            # 即使没有脱敏，也保存一个副本
            doc.save(output_pdf_path)
            doc.close()
            return

        print(f"大模型识别到以下敏感信息：{set(sensitive_keywords)}")

        # 遍历每一页，根据大模型识别出的关键词进行脱敏
        for page_num in range(len(doc)):
            page = doc.load_page(page_num)

            for keyword in set(sensitive_keywords):
                keyword_rects = page.search_for(keyword)

                for rect in keyword_rects:
                    # 使用 draw_rect 方法，用黑色矩形覆盖敏感信息
                    page.draw_rect(rect, color=(0, 0, 0), fill=(0, 0, 0))

        # 保存修改后的PDF
        doc.save(output_pdf_path)
        doc.close()
        print(f"PDF脱敏成功，已保存至：{output_pdf_path}")

    except Exception as e:
        print(f"处理PDF时发生错误: {e}")


# --- 3. 使用示例 ---
input_file = '111.pdf'
output_file = '111_redacted.pdf'

redact_pdf_with_llm(input_file, output_file)
